import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_sparse import SparseTensor

import random

class LMSPS_Se(nn.Module):

    # hidden: 隐藏层的维度。
    # nclass: 分类的类别数。
    # feat_keys: 特征键的列表。
    # label_feat_keys: 标签特征键的列表。
    # tgt_key: 目标键，通常是需要进行预测的特征。
    # dropout, input_drop: 丢弃率，用于防止过拟合。
    # device: 设备信息（如 GPU）。
    # num_final: 最终保留的路径数量。
    # residual: 是否使用残差连接。
    # bns: 是否使用批归一化。
    # data_size: 数据大小，用于初始化嵌入。
    # num_sampled: 每次采样的数量。
    def __init__(self, hidden, nclass, feat_keys, label_feat_keys, tgt_key, dropout, 
                 input_drop, device, num_final, residual=False, bns=False, data_size=None, num_sampled=1):
        
        super(LMSPS_Se, self).__init__()
        # 初始化特征和标签特征的键，以及路径的数量和样本数量等。
        self.feat_keys = feat_keys
        self.label_feat_keys = label_feat_keys
        self.num_feats = len(feat_keys)
        self.all_meta_path = list(self.feat_keys) + list(self.label_feat_keys)
        self.num_sampled = num_sampled
        self.num_channels = self.num_sampled
        self.num_paths = len(self.all_meta_path)
        self.num_final = num_final
        self.num_res = self.num_paths - self.num_final
        self.tgt_key = tgt_key
        self.residual = residual

        # 打印特征路径和标签特征路径的数量和。
        print("number of paths", len(feat_keys), len(label_feat_keys))

        # 初始化特征嵌入字典，使用均匀分布初始化权重。
        self.embeding = nn.ParameterDict({})
        for k, v in data_size.items():
            self.embeding[str(k)] = nn.Parameter(
                torch.Tensor(v, hidden).uniform_(-0.5, 0.5))

        # 初始化标签嵌入字典，使用均匀分布初始化权重。
        if len(label_feat_keys):
            self.labels_embeding = nn.ParameterDict({})
            for k in label_feat_keys:
                self.labels_embeding[k] = nn.Parameter(
                    torch.Tensor(nclass, hidden).uniform_(-0.5, 0.5))

        # 定义最终输出层，包含一个线性层和一个批归一化层。
        self.lr_output = nn.Sequential(
            nn.Linear(hidden, nclass, bias=False),
            nn.BatchNorm1d(nclass, affine=bns, track_running_stats=bns)
        )

        # 定义 PReLU 激活函数和丢弃层。
        self.prelu = nn.PReLU()
        self.dropout = nn.Dropout(dropout)
        self.input_drop = nn.Dropout(input_drop)

        # 初始化 α 参数，用于路径的权重，并允许其在训练中更新。
        self.alpha = torch.ones(self.num_paths).to(device)
        self.alpha.requires_grad_(True)

        # 如果使用残差连接，则定义一个线性层用于残差操作。
        if self.residual:
            self.res_fc = nn.Linear(hidden, hidden)

        self.init_params()

    def init_params(self):
        # 初始化权重使用 Xavier 初始化法，适用于 ReLU 激活函数。
        gain = nn.init.calculate_gain("relu")
        if self.residual:
            nn.init.xavier_uniform_(self.res_fc.weight, gain=gain)
        for layer in self.lr_output:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_uniform_(layer.weight, gain=gain)
                if layer.bias is not None:
                    nn.init.zeros_(layer.bias)

    # 返回当前的 α 参数。
    def alphas(self):
        alphas= [self.alpha]
        return alphas

    # 通过对 α 参数排序来选择需要的路径，并进行随机抽样。
    def epoch_sample(self, eps):
        # torch.argsort(self.alpha, dim=-1, descending=True)：对 α 参数进行排序，返回排序后的索引。descending=True 表示按照降序排序，权重大的在前面。
        # [:int(self.num_res * eps) + self.num_final]：从排序后的索引中选择前 (int(self.num_res * eps) + self.num_final) 个索引。这意味着根据当前的 eps（通常是一个与训练进度相关的参数），选择一部分最大的 α 值对应的路径索引。
        indices = torch.argsort(self.alpha, dim=-1, descending=True)[:int(self.num_res * eps) + self.num_final]

        # list(indices.cpu().numpy())：将选中的索引从 PyTorch 张量转换为 NumPy 数组，并进一步转换为 Python 列表。这是因为 random.sample 函数要求输入是一个列表。
        # random.sample(..., self.num_sampled)：从选中的索引中随机选择 self.num_sampled 个索引。这样可以引入随机性，避免过拟合于特定的路径选择。
        sampled = random.sample(list(indices.cpu().numpy()), self.num_sampled)
        # 对随机采样得到的索引进行排序，以便于后续处理和比较。
        sampled = sorted(sampled)
        # 输出当前采样的路径索引，方便调试和跟踪。
        print(f"sampled: {sampled}")
        return sampled
    

    # 定义前向传播过程，输入采样的路径和特征字典。
    def forward(self, epoch_sampled, feats_dict, label_feats_dict, meta_path_sampled, label_meta_path_sampled):
        # 根据输入特征的类型（稠密或稀疏张量）进行特征处理并应用嵌入。
        if isinstance(feats_dict[meta_path_sampled[-1]], torch.Tensor):
            for k, v in feats_dict.items():
                feats_dict[k] = self.input_drop(v @ self.embeding[k])
        elif isinstance(feats_dict[meta_path_sampled[-1]], SparseTensor):
            for k, v in feats_dict.items():
                feats_dict[k] = self.input_drop(v @ self.embeding[k[-1]])
        else:
            assert 0

        # 对标签特征应用嵌入和丢弃层处理。
        for k, v in label_feats_dict.items():
            label_feats_dict[k] = self.input_drop(v @ self.labels_embeding[k])

        # 实现 Gumbel 采样，用于从 logits 中获取概率分布，以便进行无偏采样。
        def get_gumbel_prob(xins):
            while True:
                # torch.empty_like(xins)：创建一个与 xins 形状相同的张量。
                # exponential_()：对这个张量进行指数分布的采样，生成符合 Gumbel 分布的随机数。
                # -...log()：对采样结果取负对数，得到 Gumbel 变量 gumbels。
                gumbels = -torch.empty_like(xins).exponential_().log()
                # xins.log_softmax(dim=-1)：对输入 xins 进行 log-softmax 操作，这样可以将 logits 转换为对数概率。
                # + gumbels：将 Gumbel 变量加到 logits 上，以引入随机性。
                # / self.tau：对结果进行温度缩放（temperature scaling），self.tau 控制采样的随机性，较低的温度会使得输出更加确定，而较高的温度则会使得输出更加随机。
                logits = (xins.log_softmax(dim=-1) + gumbels) / self.tau
                # 对 logits 进行 softmax 操作，得到最终的概率分布 probs。
                probs = nn.functional.softmax(logits, dim=-1)
                # probs.max(-1, keepdim=True)：获取最大概率的值和对应的索引。
                # [-1]：提取最大概率的索引。
                index = probs.max(-1, keepdim=True)[-1]
                # torch.zeros_like(logits)：创建一个与 logits 形状相同的零张量。
                # scatter_(-1, index, 1.0)：在最后一个维度（对应类别）将最大索引的位置设为 1，形成 one-hot 编码 one_h。
                one_h = torch.zeros_like(logits).scatter_(-1, index, 1.0)
                # one_h - probs.detach() + probs：这里的 one_h 是硬选择（hard selection），probs.detach() 用于从计算图中分离出概率。这个操作的目的是为了将 gradients 传递给 probs 而不是 one_h，从而在反向传播中不会影响到硬选择的过程。
                hardwts = one_h - probs.detach() + probs
                # 检查 gumbels、probs 是否包含无穷大或 NaN 的值。如果发现任何无穷大或 NaN，则重新进行采样（继续 while 循环）。否则，退出循环。
                if (
                        (torch.isinf(gumbels).any())
                        or (torch.isinf(probs).any())
                        or (torch.isnan(probs).any())
                ):
                    continue
                else:
                    break

            return probs

        # 将普通特征和标签特征堆叠成一个张量，形状为 [B, C, D]。
        x = [feats_dict[k] for k in meta_path_sampled] + [label_feats_dict[k] for k in label_meta_path_sampled]
        x = torch.stack(x, dim=1) # [B, C, D]

        # 获取当前采样的 α 权重并进行 Gumbel 采样。
        ws = [self.alpha[idx] for idx in epoch_sampled]
        ws = get_gumbel_prob(torch.stack(ws))

        # 使用爱因斯坦求和约定对特征进行加权。
        x = torch.einsum('bcd,c->bd', x, ws)

        # 如果使用残差连接，则将目标特征与当前输出相加。
        if self.residual:
            k = self.tgt_key
            tgt_feat = feats_dict[k]
            x = x + self.res_fc(tgt_feat)

        # 通过激活函数和丢弃层处理后，将结果传递给输出层。
        x = self.dropout(self.prelu(x))
        x = self.lr_output(x)
        
        return x

    # 用于设置 Gumbel 采样中的 tau 参数。
    def set_tau(self, tau):
        self.tau = tau

    # 采样函数，通过软最大值和阈值选择路径。
    def sample(self, keys, label_keys, lam, topn, all_path=False):
        '''
        to sample one candidate edge type per link
        '''
        length = len(self.alpha)
        seq_softmax = None if self.alpha is None else F.softmax(self.alpha, dim=-1)
        max = torch.max(seq_softmax, dim=0).values
        min = torch.min(seq_softmax, dim=0).values
        threshold = lam * max + (1 - lam) * min


        _, idxl = torch.sort(seq_softmax, descending=True)  # descending为alse，升序，为True，降序

        idx = idxl[:self.num_sampled]

        if all_path:
            path = []
            label_path = []
            for i, index in enumerate(idxl):
                if index < len(keys):
                    path.append((keys[index], i))
                else:
                    label_path.append((label_keys[index - len(keys)], i))
            return [path, label_path], idx

        if topn:
            id_paths = idxl[:topn]
        else:
            id_paths = [k for k in range(length) if seq_softmax[k].item() >= threshold]
        path = [keys[i] for i in range(len(keys)) if i in id_paths]
        label_path = [label_keys[i] for i in range(len(label_keys)) if i+len(keys) in id_paths]

        return [path, label_path], idx





